from transformers import AutoModel, AutoTokenizer, AutoModelForSequenceClassification
from transformers import pipeline
from torch.nn.functional import softmax

checkpoint = 'E:\\ai\\huggingface\\models\\distilbert-base-uncased-finetuned-sst-2-english'

"""
transformers 从 4.31.0 版本 开始支持 safetensors 格式
如果只有safetensors文件，报错Could not load model xxx with any of the following classes,可能是transformers版本太低了
"""

raw_input = ["I've been waiting for a HuggingFace course my whole life.", "I hate this so much!", "这家旅馆的床太脏了，早餐也是糟糕透顶"]
# classifier = pipeline("sentiment-analysis",
#                       model=model_path,
#                       trust_remote_code=True)
# result = classifier(raw_input)
# print(result)
print('-'*50)

tokenizer = AutoTokenizer.from_pretrained(checkpoint)
inputs = tokenizer(raw_input, padding=True, truncation=True, return_tensors='pt')
print(inputs)

model = AutoModel.from_pretrained(checkpoint)
outputs = model(**inputs)
print(outputs.last_hidden_state.shape)

model = AutoModelForSequenceClassification.from_pretrained(checkpoint)
outputs = model(**inputs)
print(outputs.logits.shape)
print(outputs.logits)

predications = softmax(outputs.logits, dim=0)
print(predications)
print(model.config.id2label)
